import re

from common import utils
from eval_lib.common.ssh import SSHPool
from common import utils as common_utils
from common.const import EB_RPM_URL
from common.utils import ssh_pool_default
from eval_lib.common.logger import get_logger

log = get_logger()


class Traffic():

    def __init__(self, traffic_ip, ssh_pool: SSHPool = ssh_pool_default):
        self.traffic_ip = traffic_ip
        self.ssh_pool = ssh_pool

    def set_traffic_commands(self, traffic_commands: list):
        self.traffic_commands = traffic_commands

    def extract_command_rate(self, command: str) -> int:
        rate = 0
        if command.strip().startswith("wrk2"):
            rate = self.extract_number_after_field("-R", command)
        elif command.strip().startswith("eb"):
            rate = self.extract_number_after_field("-r", command)
        elif command.strip().startswith("dnsperf"):
            rate = self.extract_number_after_field("-Q", command)
        elif command.strip().startswith("kubectl"):
            rate = 0
        else:
            log.error("traffic tool not support")
            rate = 0
            # assert False
        return rate

    def analyze_command(self, command: str, **kwargs) -> str:
        if command.strip().startswith("wrk2"):
            log_filter = '''|grep -E "(Latency Distribution|Requests/sec)" -A 8 | grep -E "^( 50.000| 90.000|Requests/sec:)"| awk '{print $2}'
            '''
        elif command.strip().startswith("eb"):
            log_filter = '''|awk 'END{print $14; print $16; print $8}'
            '''
        elif command.strip().startswith("dnsperf"):
            log_filter = '''|awk '/Queries per second/{rps=$4} /Average Latency/{avg_latecy=$4; max_latecy=substr($8, 1, length($8)-1)} END{print avg_latecy"s"; print max_latecy"s"; print rps}'
            '''
        elif command.strip().startswith("kubectl"):
            log_filter = ''
        else:
            log.error("traffic tool not support")
            log_filter = ''
            # assert False
        for key, value in kwargs.items():
            command = command.replace(f"${key}", str(value))
        rel_cmd = command + log_filter
        return rel_cmd

    def get_traffic_tool_result(self, command: str, data: list) -> dict:
        cmd = command.strip()
        rate = self.extract_command_rate(cmd)
        result = {}
        if cmd.startswith("wrk2") or cmd.startswith("eb") or cmd.startswith(
                "dnsperf"):
            result["server.latency_p50"] = utils.format_latency(
                data[0].strip(), "ms")
            result["server.latency_p90"] = utils.format_latency(
                data[1].strip(), "ms")
            result["server.rps"] = data[2].strip()
            result["server.rate"] = rate
            result["case.command"] = cmd
        elif cmd.startswith("kubectl"):
            pass
        return result

    def generate_traffic(self, traffic_command, **kwargs):
        ssh = self.ssh_pool.get(self.traffic_ip)
        cmd = self.analyze_command(traffic_command, **kwargs)
        log.info(f"traffic command: {cmd}")
        _, stdout, stderr = ssh.exec_command(cmd)
        output = stdout.readlines()
        error = stderr.readlines()
        exit_status = stdout.channel.recv_exit_status()
        if exit_status == 0:
            log.info("traffic generate success")
            result = self.get_traffic_tool_result(traffic_command, output)
        else:
            log.error(f"traffic generate failed, err: {error}")
            result = {}
        return result

    def install_eb_rpm(self, vm_ip=None):
        if not vm_ip:
            vm_ip = self.traffic_ip
        if common_utils.check_ssh_command(self.traffic_ip, "eb --help",
                                          self.ssh_pool):
            log.info("eb is already installed")
            return
        log.info("start to install eb")
        ssh_client = self.ssh_pool.get(vm_ip)
        _, stdout, stderr = ssh_client.exec_command(
            f''' curl -O {EB_RPM_URL} && \
            rpm -ivh eval-bench-1.0-* ''')
        output = stdout.read().decode()
        error = stderr.read().decode()
        if "eval-bench" in output:
            log.info("Install eval-bench successfully")
        else:
            log.error(f"Install eval-bench failed, error is {error}")
            assert False

    def extract_number_after_field(self, field, text):
        pattern = r"{}\s*(\d+)".format(field)
        # 在文本开头搜索匹配的内容
        match = re.search(pattern, text)
        if match:
            # 如果找到匹配的内容，返回匹配的数字部分
            return match.group(1)
        else:
            # 如果没有找到匹配的内容，返回 None
            return None

    def create_performance_pods(self, vm_ip):
        '''Login to the vtaps by SSH,create pod for UDP performance
        '''
        utils.set_nameserver(vm_ip, self.ssh_pool)
        ssh_client = self.ssh_pool.get(vm_ip)
        cmds = [
            "curl -O http://nexus.yunshan.net/repository/tools/automation/performance/performance_client.yaml",
            "curl -O http://nexus.yunshan.net/repository/tools/automation/performance/performance_server.yaml",
            "sed -i 's/image: performance:lastest/image: dfcloud-image-registry.cn-beijing.cr.aliyuncs.com\/public\/agent_performance_test:latest/g' performance_client.yaml",
            "sed -i 's/image: performance:lastest/image: dfcloud-image-registry.cn-beijing.cr.aliyuncs.com\/public\/agent_performance_test:latest/g' performance_server.yaml",
            "kubectl apply -f performance_client.yaml -f performance_server.yaml ",
            "kubectl wait --for=condition=Ready pod -l app=performance-client-pod --timeout=300s",
            "kubectl wait --for=condition=Ready pod -l app=performance-server-pod --timeout=300s",
        ]
        cmd = " && ".join(cmds)
        log.info("start create perfomance pods")
        _, stdout, _ = ssh_client.exec_command(cmd)
        logs = stdout.readlines()
        log.info(f"log: {logs}")
        pod_client_name = ''
        pod_client_ip = ''
        pod_server_name = ''
        pod_server_ip = ''
        _, stdout, _ = ssh_client.exec_command(
            "kubectl get pods -o wide | grep -E 'client|server'")
        pods_info = stdout.readlines()
        log.info(f"pods info: {pods_info}")
        for i in pods_info:
            if 'client' in i:
                pod_client_name = i.split()[0]
                pod_client_ip = i.split()[5]
            elif 'server' in i:
                pod_server_name = i.split()[0]
                pod_server_ip = i.split()[5]
        return pod_client_name, pod_client_ip, pod_server_name, pod_server_ip
